import sensor,image,lcd  # import 相关库
import utime
import time
import gc
import uos
import KPU as kpu
import uart
import gpio

task_fd = None
task_ld = None
task_fe = None
anchor = (1.889, 2.5245, 2.9465, 3.94056, 3.99987, 5.3658, 5.155437, 6.92275, 6.718375, 9.01025) #anchor for face detect 用于人脸检测的Anchor
# anchorMsk = (0.1606, 0.3562, 0.4712, 0.9568, 0.9877, 1.9108, 1.8761, 3.5310, 3.4423, 5.6823)
dst_point = [(44,59),(84,59),(64,82),(47,105),(81,105)] #standard face key point position 标准正脸的5关键点坐标 分别为 左眼 右眼 鼻子 左嘴角 右嘴角
img_lcd=image.Image() # 设置显示buf
img_face=image.Image(size=(128,128)) #设置 128 * 128 人脸图片buf
a=img_face.pix_to_ai() # 将图片转为kpu接受的格式
record_ftr=[] #空列表 用于存储当前196维特征

#空列表 用于存储按键记录下人脸特征， <<<!可以将特征以txt等文件形式保存到sd卡后，读取到此列表，即可实现人脸断电存储!
record_ftrs=[]
facial_data = {}
name = ''
addIDFlag = 0
delIDFlag = 0

def faceInit():
    global task_fd
    global task_ld
    global task_fe
    global facial_data
    task_fd = kpu.load('/sd/model/FD_1d9dcdb3eaad35d3bb85d3aefa482c89.smodel') # 从sd卡中加载人脸检测模型
    task_ld = kpu.load('/sd/model/KP_chwise_1d9dcdb3eaad35d3bb85d3aefa482c89.smodel') # 从sd卡中加载人脸五点关键点检测模型
    task_fe = kpu.load('/sd/model/FE_mbv1_0.5_1d9dcdb3eaad35d3bb85d3aefa482c89.smodel') # 从sd卡中加载人脸196维特征值模型
    a = kpu.init_yolo2(task_fd, 0.5, 0.3, 5, anchor) #初始化人脸检测模型
    ftrFileList = uos.listdir('/sd/facial_data')
    for x in ftrFileList:
        fileName = x[:-4]
        fobj = open('/sd/facial_data/' + x, 'rb')
        facial_data[fileName] = fobj.read()
        fobj.close()
    print(facial_data)

def faceDeInit():
    a = kpu.deinit(task_fe)
    a = kpu.deinit(task_ld)
    a = kpu.deinit(task_fd)

def faceRecog(img):
    global addIDFlag
    global delIDFlag
    code = kpu.run_yolo2(task_fd, img) # 运行人脸检测模型，获取人脸坐标位置
    tmpName = uart.uart2.read()
    if tmpName:
        try:
            name = tmpName.decode('utf-8')
            if 'add:' in name:
                name = name[4:]
                addIDFlag = 1
            elif 'del:' in name:
                name = name[4:]
                delIDFlag = 1
            print(name)
        except:
            print('请确保输入是utf-8的汉字')
        tmpName = None
    
    
    if code: # 如果检测到人脸
        #face#VECTOR += 1
        for i in code: # 迭代坐标框
            # Cut face and resize to 128x128
            a = img.draw_rectangle(i.rect()) # 在屏幕显示人脸方框
            face_cut=img.cut(i.x(),i.y(),i.w(),i.h()) # 裁剪人脸部分图片到 face_cut
            face_cut_128=face_cut.resize(128,128) # 将裁出的人脸图片 缩放到128 * 128像素
            a=face_cut_128.pix_to_ai() # 将裁出图片转换为kpu接受的格式
            #a = img.draw_image(face_cut_128, (0,0))
            # Landmark for face 5 points
            fmap = kpu.forward(task_ld, face_cut_128) # 运行人脸5点关键点检测模型
            plist=fmap[:] # 获取关键点预测结果
            le=(i.x()+int(plist[0]*i.w() - 10), i.y()+int(plist[1]*i.h())) # 计算左眼位置， 这里在w方向-10 用来补偿模型转换带来的精度损失
            re=(i.x()+int(plist[2]*i.w()), i.y()+int(plist[3]*i.h())) # 计算右眼位置
            nose=(i.x()+int(plist[4]*i.w()), i.y()+int(plist[5]*i.h())) #计算鼻子位置
            lm=(i.x()+int(plist[6]*i.w()), i.y()+int(plist[7]*i.h())) #计算左嘴角位置
            rm=(i.x()+int(plist[8]*i.w()), i.y()+int(plist[9]*i.h())) #右嘴角位置
            #a = img.draw_circle(le[0], le[1], 4)
            #a = img.draw_circle(re[0], re[1], 4)
            #a = img.draw_circle(nose[0], nose[1], 4)
            #a = img.draw_circle(lm[0], lm[1], 4)
            #a = img.draw_circle(rm[0], rm[1], 4) # 在相应位置处画小圆圈
            # align face to standard position
            src_point = [le, re, nose, lm, rm] # 图片中 5 坐标的位置
            T=image.get_affine_transform(src_point, dst_point) # 根据获得的5点坐标与标准正脸坐标获取仿射变换矩阵
            a=image.warp_affine_ai(img, img_face, T) #对原始图片人脸图片进行仿射变换，变换为正脸图像
            a=img_face.ai_to_pix() # 将正脸图像转为kpu格式
            #a = img.draw_image(img_face, (128,0)) #选择不显示仿射变换后的图像
            del(face_cut_128) # 释放裁剪人脸部分图片
            # calculate face feature vector
            fmap = kpu.forward(task_fe, img_face) # 计算正脸图片的196维特征值
            feature=kpu.face_encode(fmap[:]) #获取计算结果
            reg_flag = False
            scores = {} # 存储特征比对分数
            for j in facial_data.keys(): #迭代已存特征值
                score = kpu.face_compare(facial_data[j], feature) #计算当前人脸特征值与已存特征值的分数
                scores[j] = score #添加分数总表
            max_score = 0
            index = 0
            for k in facial_data.keys(): #迭代所有比对分数，找到最大分数和索引值
                if max_score < scores[k]:
                    max_score = scores[k]
                    index = k
            if max_score > 85: # 如果最大分数大于85， 可以被认定为同一个人
                a = img.draw_string(i.x(),i.y(), ("%s :%2.1f" % (index, max_score)), color=(0,255,0),scale=2) # 显示人名 与 分数
            else:
                a = img.draw_string(i.x(),i.y(), ("Illegal Break in! :%2.1f" % (max_score)), color=(255,0,0),scale=2) #显示未知 与 分数
                gpio.ID_out.value(1)
            if addIDFlag == 1: #如果检测到按键
                addIDFlag = 0 #重置按键状态
                record_ftr = feature
                if name:
                    facial_data[name] = record_ftr
                    fobj = open('/sd/facial_data/' + name + '.hex', 'wb')
                    fobj.write(facial_data[name])
                    fobj.close()
            break
    if delIDFlag == 1: #如果检测到按键
                delIDFlag = 0 #重置按键状态
                if name:
                    try:
                        uos.remove('/sd/facial_data/' + name + '.hex')
                        facial_data.pop(name)
                        print('del', name)
                    except:
                        print('No such file')
                    name = None

    #fps =clock.fps() #计算帧率
    #print("%2.1f fps"%fps) #打印帧率
    a = lcd.display(img) #刷屏显示
    #kpu.memtest()
    #keyAddID.disirq()
    gc.collect()